/*
============================================================================
 Name        : nmsJpegMcuReader.c
 Author      : Jian
 Version     :
 Copyright   : base on http://ce.sharif.edu/~ahmadinejad/jpeg/
 Description : nmsJpegMcuReader.c - source file
============================================================================
*/

/* INCLUDE FILES */
#include "nmsJpegMcuReader.h"
#include "nmsJpeg.h"
#include "nmsHuffman.h"
#include "nmsFrame.h"

extern const UINT8 _zigZagScanner[64] ;

nmsRgbMatrix nmsRgbCache   = {0} ;

static int _yMatrix[4][64] = {0};
static int _cbMaTrix[64]   = {0} ;
static int _crMatrix[64]   = {0} ;
static nmsMatrix* _matrixArray[__NMC_JPEG_COMP_ID_CR + 1][4] = {
    NULL, NULL, NULL, NULL,
    &(_yMatrix[0]), &(_yMatrix[1]), &(_yMatrix[2]), &(_yMatrix[3]),
    &_cbMaTrix, &_cbMaTrix, &_cbMaTrix, &_cbMaTrix,
    &_crMatrix, &_crMatrix, &_crMatrix, &_crMatrix
} ;

static UINT8 _cbCrIndexArray[64] = {0} ;
static UINT8 _cbStartIndex[4]    = {0} ;


static __NMS_INLINE int _readByte(nmsJpegMcu* mcu, UINT32* bytePos, int i)
{
    UINT8  tmpByte = 0 ;
    if ((tmpByte = mcu->fileBuf[*bytePos])== 0xFF) {
        mcu->ffPos[i] = (*bytePos + 1) * 8 ;
        *bytePos += 2 ;
        if (mcu->fileBuf[*bytePos - 1] != 0)
            nprintf("mcu _readByte error got a maker = %d", mcu->fileBuf[*bytePos - 1]) ;
    } else {
        *bytePos += 1 ;
        mcu->ffPos[i] = -1 ;
    }

    return tmpByte ;
}

static __NMS_INLINE int _read2Byte(nmsJpegMcu* mcu)
{
    UINT32 bytePos = 0 ;
    UINT32 bitPos  = 0 ;
    UINT8  tmpByte = 0 ;

    bytePos = mcu->readBitlen / 8 ;
    bitPos  = mcu->readBitlen % 8 ;

    mcu->curData = 0 ;
    mcu->curData = mcu->curData | (UINT16)(((UINT16)_readByte(mcu, &bytePos, 0)) << (8 + bitPos)) ;
    mcu->curData = mcu->curData | (UINT16)(((UINT16)_readByte(mcu, &bytePos, 1)) << bitPos) ;
    mcu->curData = mcu->curData | (UINT16)(((UINT16)_readByte(mcu, &bytePos, 2)) >> (8 - bitPos)) ;

    return 0 ;
}

static int _readFromFile(nmsJpegMcu* mcu)
{
    int readByte = 0 ;

    nmsFsSeek(mcu->jpg->fileHandle, mcu->curFilePos, NMS_FS_FILE_BEGIN) ;

    if ((readByte = nmsJpegReadFile(mcu->jpg->fileHandle, mcu->fileBuf, __NMS_JPGE_MCU_BUF_LEN)) < 0)
        return -1 ;

    mcu->bufBitLen = readByte * 8 ;

    return readByte ;
}

static __NMS_INLINE int _readDU(nmsJpegMcu* mcu, int compId)
{
    int    val       = 0 ;
    UINT8  zeroCount = 0 ;
    int i            = 1 ;
    int pRet         = 0 ;

    nmsJpeg* jpg = mcu->jpg ;

    if ((pRet = nmsJpegHuffmanParse(jpg->dcTable[jpg->sosComp[compId].dcHTIndex], mcu, 0, &val, &zeroCount)) < 0) {
        nprintf("nmsJpegHuffmanParse error") ;
        return -1 ;
    }

    mcu->preDC[compId] = mcu->preDC[compId] + val ;

    while(i < 64) {
        if ((pRet = nmsJpegHuffmanParse(jpg->acTable[jpg->sosComp[compId].acHTIndex], mcu, 1, &val, &zeroCount)) < 0) {
            nprintf("nmsJpegHuffmanParse error") ;
            return -1 ;
        }

        if (pRet == __EOB_FLAG) break ;

        i += (zeroCount + 1) ;
    }

    return 0 ;
}

static __NMS_INLINE int _readDUData(nmsJpegMcu* mcu, int compId, nmsMatrix* matrix)
{
    int    val       = 0 ;
    UINT8  zeroCount = 0 ;
    int i            = 1 ;
    int pRet         = 0 ;
    nmsMatrix tmpM   = {0} ;

    nmsJpeg*   jpg   = mcu->jpg ;
    nmsJpegDqt* qt   = jpg->qt[jpg->frame->comp[compId].qtIndex] ;

    if ((pRet = nmsJpegHuffmanParse(jpg->dcTable[jpg->sosComp[compId].dcHTIndex], mcu, 0, &val, &zeroCount)) < 0) {
        nprintf("nmsJpegHuffmanParse error") ;
        return -1 ;
    }

    tmpM[0]            = mcu->preDC[compId] + val ;
    mcu->preDC[compId] = tmpM[0] ;

    while(i < 64) {
        if ((pRet = nmsJpegHuffmanParse(jpg->acTable[jpg->sosComp[compId].acHTIndex], mcu, 1, &val, &zeroCount)) < 0) {
            nprintf("nmsJpegHuffmanParse error") ;
            return -1 ;
        }

        if (pRet == __EOB_FLAG) {
            for (; i < 64; i++) {
                tmpM[_zigZagScanner[i]] = 0 ;
            }
            break ;
        }

        if (zeroCount != 0) {
            int end = i + zeroCount ;
            for (; i < end; i++) {
                tmpM[_zigZagScanner[i]] = 0 ;
            }
        }

        tmpM[_zigZagScanner[i]] = val ;
        i += 1 ;
    }

    /* DeQuantization */
    for (i = 0; i < 64; i++ )
        tmpM[i] = tmpM[i] * (*qt)[i] ;

    /* IDCT */
    nmsJpegIdct(tmpM, matrix) ;

    return 0 ;
}

static __NMS_INLINE void _checkMaker(nmsJpegMcu* mcu)
{
    UINT32 bytePos   = 0 ;
    UINT32 bitPos    = 0 ;
    UINT8  curByte   = 0 ;
    if ((mcu->curData >> 8) != 0xFF)
        return ;

    bitPos = mcu->readBitlen % 8 ;
    bytePos = mcu->readBitlen / 8 + ((bitPos == 0) ? 0 : 1);

    if (mcu->fileBuf[bytePos] != 0xFF)
        return ;

    bytePos++ ;
    mcu->readBitlen += (8 + ((bitPos == 0) ? 0 :(8 - bitPos))) ;

    while(mcu->readBitlen < mcu->bufBitLen && mcu->fileBuf[bytePos] == 0xFF) {
        if (mcu->bufBitLen == __NMS_JPGE_MCU_BUF_LEN * 8 && (mcu->readBitlen > (mcu->bufBitLen - 8 * 6))) {
            mcu->curFilePos += bytePos ;
            _readFromFile(mcu) ;
            bytePos = 0 ;
            mcu->readBitlen = 0 ;
        }
        bytePos++ ;
        mcu->readBitlen += 8 ;
    }

    curByte = mcu->fileBuf[bytePos] ;

    nprintf("_checkMaker got a maker = %d", curByte) ;

    if (curByte == 0x00) {
        mcu->readBitlen -= 8 ; /* Jump back to FF 00 */
    } else if (curByte == 0xFF || curByte == 0xD9) {
        nprintf("_checkMaker got end of file ") ;
    } else if (curByte >= 0xD0 && curByte<=0xD7) {
        mcu->preDC[0] = 0 ;
        mcu->preDC[1] = 0 ;
        mcu->preDC[2] = 0 ;
        mcu->preDC[3] = 0 ;
        mcu->readBitlen += 8 ;
    } else {
        int len = ((UINT16)mcu->fileBuf[bytePos]) || (((UINT16)mcu->fileBuf[bytePos+1]) << 8) ;
        nprintf("_checkMaker got unknown maker len = %d", len) ;
        mcu->curFilePos += (bytePos + 2 + len + 1);
        _readFromFile(mcu) ;
        mcu->readBitlen = 0 ;
    }

    _read2Byte(mcu) ;
}


/*
* Standard
* R = Y + 1.402 *(Cr-128)
* G = Y - 0.34414*(Cb-128) - 0.71414*(Cr-128)
* B = Y + 1.772 *(Cb-128)
*/

static __NMS_INLINE UINT8 _round(int rgb)
{
    rgb += 128 ;

    if (rgb < 0) return 0 ;
    else if (rgb > 255) return 255 ;

    return (UINT8)rgb ;
}

static __NMS_INLINE void _yuvToRgbCache(nmsJpegMcu* mcu)
{
    nmsJpeg* jpg      = mcu->jpg ;
    UINT8 r           = 0 ;
    UINT8 g           = 0 ;
    UINT8 b           = 0 ;
    int cbcStartIndex = 0 ;
    int cbCrIndex     = 0 ;

    if (jpg->frame->cmpCount == __NMS_JPEG_YUV) {
        int i, j ;
        for (i = 0; i < jpg->frame->comp[__NMC_JPEG_COMP_ID_Y].duCount; i++) {
            cbcStartIndex = _cbStartIndex[i] ;
            for (j = 0; j < 64; j++) {
                cbCrIndex = _cbCrIndexArray[j] + cbcStartIndex ;
                r = _round(_yMatrix[i][j] + (1402 * _crMatrix[cbCrIndex] + 500) / 1000) ;
                g = _round(_yMatrix[i][j]  - (34414 * _cbMaTrix[cbCrIndex] - 25000) / 100000 - (71414 * _crMatrix[cbCrIndex] - 25000) / 100000) ;
                b = _round(_yMatrix[i][j] + (1772 *_cbMaTrix[cbCrIndex] + 500)/ 1000) ;
                nmsRgbCache[i][j] = (((UINT32)r) << 16) | (((UINT32)g) << 8) | ((UINT32)b) ;
            }
        }

    } else {
        /* gray picture */
        int i = 0 ;
        for (; i < 64; i++) {
            nmsRgbCache[0][i] = _yMatrix[0][i];
        }
    }
}

static int _init(nmsJpegMcu* mcu, int startBitPos, int* curLineDc)
{
    mcu->curFilePos = startBitPos / 8 ;
    mcu->readBitlen = startBitPos % 8 ;

    if (_readFromFile(mcu) <= 0) {
        nprintf("mcu reader read file error") ;
        return -1 ;
    }

    nmsMemCpy(mcu->preDC, curLineDc, sizeof(mcu->preDC)) ;

    mcu->ffPos[0] = -1 ;
    mcu->ffPos[1] = -1 ;
    mcu->ffPos[2] = -1 ;

    _read2Byte(mcu) ;

#ifdef __NMS_JPEG_DEBUG
    mcu->testReadBitTotal = 0 ;
#endif

    return 0 ;
}


static int _findNextLineMCUPos(nmsJpegMcu* mcu, int startBitPos, int* curLineDc, int* nextLineDC)
{
    int i, j, k ;

    nmsJpeg* jpg = mcu->jpg ;

    if (_init(mcu, startBitPos, curLineDc) < 0)
        return -1 ;

    for (i =0; i < jpg->mcuPerLine; i++) {
        /* Read Per-MCU without data */
        for (j = __NMC_JPEG_COMP_ID_Y; j <= __NMC_JPEG_COMP_ID_CR; j++) {
            for (k = 0; k < jpg->frame->comp[j].duCount; k++) {
                if (_readDU(mcu, j))    /* Read per du */
                    return -1 ;
            }
        }

        _checkMaker(mcu) ;
    }

    nmsMemCpy(nextLineDC, mcu->preDC, sizeof(mcu->preDC)) ;

    return mcu->curFilePos * 8 + mcu->readBitlen ;
}


nmsJpegMcu* nmsMcuCreate(struct __nms_jpeg* jpg)
{
    nmsJpegMcu* mcu = (nmsJpegMcu*)nmsJpegMalloc(sizeof(nmsJpegMcu)) ;
    memset(mcu, 0, sizeof(nmsJpegMcu)) ;
    mcu->jpg = jpg ;
    return mcu ;
}

void nmsMcuDesctroy(nmsJpegMcu* mcu)
{
    if (mcu) {
        nmsJpegFree(mcu) ;
    }
}

int nmsMcuInit(nmsJpegMcu* mcu, int startBitPos, int mcuVLineNo, int* curLineDc, int* nextLineDC)
{
    int nextPos = -1 ;

    __NMS_ASSERT_DEFENSE(mcu, return -1) ;

    if (nextLineDC && (nextPos = _findNextLineMCUPos(mcu, startBitPos, curLineDc, nextLineDC)) < 0) {
        nprintf("nmsMcuInit find next line mcu pos error") ;
        return -1 ;
    }

    _init(mcu, startBitPos,curLineDc) ;

    mcu->x = -mcu->jpg->mcuWidth ;
    mcu->y = mcuVLineNo * mcu->jpg->mcuHeight ;

    mcu->lineStartBitPos = startBitPos ;
    mcu->readDUCount     = 0 ;

    return nextPos ;
}

int nmsMcuEatBits(nmsJpegMcu* mcu, int eatBitLen)
{
#ifdef __NMS_JPEG_DEBUG
    mcu->testReadBitTotal += eatBitLen ;
#endif

    mcu->readBitlen += eatBitLen ;

    if (mcu->readBitlen >= mcu->ffPos[0])
        mcu->readBitlen += 8 ;

    if (mcu->readBitlen >= mcu->ffPos[1])
        mcu->readBitlen += 8 ;

    if (mcu->readBitlen >= mcu->ffPos[2])
        mcu->readBitlen += 8 ;

    __NMS_ASSERT_DEFENSE(mcu->readBitlen <= mcu->bufBitLen, return -1) ;

    /* 8 * 6 mean that we may got a data with FF 00 FF 00 FF 00 */
    if (mcu->bufBitLen == __NMS_JPGE_MCU_BUF_LEN * 8 && (mcu->readBitlen > (mcu->bufBitLen - 8 * 6))) {
        mcu->curFilePos += mcu->readBitlen / 8  ;
        mcu->readBitlen = mcu->readBitlen % 8 ;
        _readFromFile(mcu) ;
    }

    _read2Byte(mcu) ;

    return 0 ;
}


int nmsMcuReadNext(nmsJpegMcu* mcu)
{
    int k, l;

    nmsJpeg* jpg = mcu->jpg ;

    mcu->readDUCount++ ;

    __NMS_ASSERT(mcu->readDUCount <= mcu->jpg->mcuPerLine) ;

    mcu->x += mcu->jpg->mcuWidth ;

    /* Read Per-MCU with data */
    for (k = __NMC_JPEG_COMP_ID_Y; k <= __NMC_JPEG_COMP_ID_CR; k++) {
        for (l = 0; l < jpg->frame->comp[k].duCount; l++) {
            if (_readDUData(mcu, k, _matrixArray[k][l]))    /* Read per-du */
                return -1 ;
        }
    }

    /* Start to Parser YUV to RGB */
    _yuvToRgbCache(mcu) ;

    /* End of MCU, we should check if there a maker */
    _checkMaker(mcu) ;

    return 0 ;
}


int nmsMcuCmpXPixel(const nmsJpegMcu* mcu, int x)
{
    __NMS_ASSERT_DEFENSE(mcu, return 0) ;
    if (x < mcu->x)                         return -1 ;
    if (x >= (mcu->x + mcu->jpg->mcuWidth)) return 1 ;
    return 0 ;
}

int nmsMcuCmpYPixel(const nmsJpegMcu* mcu, int y)
{
    __NMS_ASSERT_DEFENSE(mcu, return 0) ;
    if (y < mcu->y)                          return -1 ;
    if (y >= (mcu->y + mcu->jpg->mcuHeight)) return 1;
    return 0 ;
}

int nmsMcuGetNextLineStartPos(nmsJpegMcu* mcu)
{
    int k, j ;

    while(mcu->readDUCount < mcu->jpg->mcuPerLine) {
        for (j = __NMC_JPEG_COMP_ID_Y; j <= __NMC_JPEG_COMP_ID_CR; j++) {
            for (k = 0; k < mcu->jpg->frame->comp[j].duCount; k++) {
                if (_readDU(mcu, j))    /* Read per du */
                    return -1 ;
            }
        }
        _checkMaker(mcu) ;
        mcu->readDUCount++ ;
    }

    return mcu->curFilePos * 8 + mcu->readBitlen ;
}


void nmsMcuInitCbCrIndexArray(struct __nms_jpeg* jpg)
{
    int i = 0 ;

    nmsMemSet(_cbStartIndex, 0, sizeof(_cbStartIndex)) ;

    if (jpg->frame->comp[__NMC_JPEG_COMP_ID_Y].duCount == 4) {
        for (; i < 64; i++) {
            _cbCrIndexArray[i] = ((i % 8) / 2) + ((i / 16) * 8)  ;
        }
        _cbStartIndex[0] = 0 ;
        _cbStartIndex[1] = 4 ;
        _cbStartIndex[2] = 32 ;
        _cbStartIndex[3] = 36 ;
    } else if (jpg->frame->comp[__NMC_JPEG_COMP_ID_Y].duCount == 2) {
        if (jpg->frame->comp[__NMC_JPEG_COMP_ID_Y].hSample == 2) {
            for (; i < 64; i++) {
                _cbCrIndexArray[i] = ((i % 8) / 2) + ((i / 8) * 8) ;
            }
            _cbStartIndex[0] = 0 ;
            _cbStartIndex[1] = 4 ;
        } else {
            for (; i < 64; i++) {
                _cbCrIndexArray[i] = (i % 8) + ((i / 16) * 8);
            }
            _cbStartIndex[0] = 0 ;
            _cbStartIndex[1] = 32 ;
        }
    } else {

        for (; i < 64; i++) {
            _cbCrIndexArray[i] = i ;
        }
    }
}


#ifdef __NMS_JPEG_DEBUG

int nmsMcuTestReadDU(nmsJpegMcu* mcu, int compId, nmsMatrix* oriMatrix, nmsMatrix* dqtMatrix, nmsMatrix* idctMatrix)
{
    /* a copy of _readDUData just hide some static data */

    int    val       = 0 ;
    UINT8  zeroCount = 0 ;
    int i            = 1 ;
    int pRet         = 0 ;

    nmsJpeg*   jpg   = mcu->jpg ;
    nmsJpegDqt* qt   = jpg->qt[jpg->frame->comp[compId].qtIndex] ;

    if ((pRet = nmsJpegHuffmanParse(jpg->dcTable[jpg->sosComp[compId].dcHTIndex], mcu, 0, &val, &zeroCount)) < 0) {
        nprintf("nmsJpegHuffmanParse error") ;
        return -1 ;
    }

    (*oriMatrix)[0]    = mcu->preDC[compId] + val ;
    mcu->preDC[compId] = (*oriMatrix)[0] ;
    (*dqtMatrix)[0]    = (*oriMatrix)[0] ;

    while(i < 64) {
        if ((pRet = nmsJpegHuffmanParse(jpg->acTable[jpg->sosComp[compId].acHTIndex], mcu, 1, &val, &zeroCount)) < 0) {
            nprintf("nmsJpegHuffmanParse error") ;
            return -1 ;
        }

        if (pRet == __EOB_FLAG) {
            for (; i < 64; i++) {
                (*dqtMatrix)[_zigZagScanner[i]] = 0 ;
                (*oriMatrix)[i]                 = 0 ;
            }
            break ;
        }

        if (zeroCount != 0) {
            int end = i + zeroCount ;
            for (; i < end; i++) {
                (*dqtMatrix)[_zigZagScanner[i]] = 0 ;
                (*oriMatrix)[i]                 = 0 ;
            }
        }

        (*dqtMatrix)[_zigZagScanner[i]] = val ;
        (*oriMatrix)[i]                 = val ;
        i += 1 ;
    }

    /* DeQuantization */
    for (i = 0; i < 64; i++ )
        (*dqtMatrix)[i] = (*dqtMatrix)[i] * (*qt)[i] ;

    /* IDCT */
    nmsJpegIdct(dqtMatrix, idctMatrix) ;

    return 0 ;
}

void nmsMcuTestCheckMaker(nmsJpegMcu* mcu)
{
    _checkMaker(mcu) ;
}

#endif


/* End of File */
